package com.artislong.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * @author 陈敏
 * @version AgentConfiguration.java, v 1.0 2025 07 22 14:45 chenmin Exp $
 * Created on 2025 07 22
 */
@Configuration
public class AgentConfiguration {

    @Bean
    public ChatClient chatClient(ChatClient.Builder builder, ChatMemory chatMemory, VectorStore vectorStore, Advisor documentRetrieverAdvisor) {
        return builder
                .defaultAdvisors(
                        SimpleLoggerAdvisor.builder().build(),
                        // RAG
                        QuestionAnswerAdvisor.builder(vectorStore).build(),
                        // 文档检索
                        documentRetrieverAdvisor,
                        // 会话存储
                        MessageChatMemoryAdvisor.builder(chatMemory).build()
                )
                .build();
    }

    /**
     *
     * @param retriever
     * @return
     */
    @Bean
    public Advisor documentRetrieverAdvisor(DocumentRetriever retriever) {
        return RetrievalAugmentationAdvisor.builder()
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)
                        .build())
                .documentRetriever(retriever)
                .build();
    }

    /**
     * 配置文档检索器
     * @param vectorStore
     * @return
     */
    @Bean
    public DocumentRetriever documentRetriever(VectorStore vectorStore) {
        return VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.5)    // 设置相似度阈值
                .topK(3)                     // 返回前3个最相关的文档
                .build();
    }
}
